import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import minimize
from tools.Cal_Z import Cal_Z
from tools.read_hdf5_find_powers import read_hdf5_find_powers
from tools.mag_phase_to_IQ import mag_phase_to_IQ
from tools.calculate_deri import calculate_deri
from tools.Fdecay_fhalf import Fdecay_fhalf

# 系统参数配置
class Config:
    RQ = 8893  # 电阻-品质因数乘积
    Ql = 76.7  # 外部品质因数
    beta = 1.2
    Fs = 117.402e6 / 14  # 采样率 [Hz]
    Ts = 1 / Fs  # 采样时间间隔 [s]

# 数据源配置
DATA_SOURCES = [
    r'D:\Beam phase calibration\MEBT_50us_scan_beam_current_3mA_63mA_B1_B2\AFF\63mA_70mm_aff\RFQ-Bunchers_3-60mA_50us_1628_2023May25.hdf5'
]

class VoltageProcessor:
    @staticmethod
    def clean_voltages(voltages_IQ):
        """清除电压数据中的无效值（NaN或无穷大）"""
        return {key: np.nan_to_num(value) for key, value in voltages_IQ.items()}

    @staticmethod
    def calculate_calibration_matrix(voltages_IQ):
        """计算校准矩阵和校准后的电压值"""
        voltages_IQ = VoltageProcessor.clean_voltages(voltages_IQ)
        
        # 提取IQ分量
        Vcav_I, Vcav_Q = voltages_IQ["Vcav_I"], voltages_IQ["Vcav_Q"]
        Vfor_I, Vfor_Q = voltages_IQ["Vfor_I"], voltages_IQ["Vfor_Q"]
        Vrefl_I, Vrefl_Q = voltages_IQ["Vrefl_I"], voltages_IQ["Vrefl_Q"]
        
        # 构建校准矩阵
        A_matrix = np.column_stack((
            np.append(Vfor_I, Vfor_Q),
            np.append(-Vfor_Q, Vfor_I),
            np.append(Vrefl_I, Vrefl_Q),
            np.append(-Vrefl_Q, Vrefl_I)
        ))
        y_array = np.append(Vcav_I, Vcav_Q)

        try:
            C_array = np.matmul(np.linalg.pinv(A_matrix, rcond=1e-5), y_array.transpose())
        except np.linalg.LinAlgError as e:
            print(f"SVD计算失败: {e}")
            return None

        # 计算校准后的电压值
        Vfor_cali = (Vfor_I + 1j * Vfor_Q) * (C_array[0] + 1j * C_array[1])
        Vrefl_cali = (Vrefl_I + 1j * Vrefl_Q) * (C_array[2] + 1j * C_array[3])

        return {
            'Vfor_I_cali': Vfor_cali.real,
            'Vfor_Q_cali': Vfor_cali.imag,
            'Vrefl_I_cali': Vrefl_cali.real,
            'Vrefl_Q_cali': Vrefl_cali.imag,
            'C_array': C_array
        }

class OptimizationCalculator:
    @staticmethod
    def calculate_y_a(a, Z, Vc, Vf, Vr):
        """计算y_a参数"""
        w_half = Fdecay_fhalf(Vc, Vf, Config.Fs)
        a_complex = complex(a[0], a[1])
        Vf_cal = a_complex * Vf + a_complex / Z * Vr
        rho = 2 * np.abs((Config.beta / (1 + Config.beta)) * Vf_cal)
        return w_half * (rho * np.cos(np.angle(Vc) - np.angle(Vf_cal)) - np.abs(Vc))

    @staticmethod
    def error_function(a, Z, Vc, Vf, Vr):
        """计算误差函数"""
        dVc_dt = calculate_deri(np.abs(Vc), Config.Fs)
        y_a = OptimizationCalculator.calculate_y_a(a, Z, Vc, Vf, Vr)
        return np.sum((dVc_dt - y_a) ** 2)

    @staticmethod
    def optimize_parameters(Vc, Vf, Vr, Z):
        """优化a参数"""
        result = minimize(
            OptimizationCalculator.error_function,
            [0, 0.0],
            args=(Z, Vc, Vf, Vr),
            method='Nelder-Mead'
        )
        print(f"最小误差: {result.fun}")
        return complex(result.x[0], result.x[1])

class Plotter:
    @staticmethod
    def plot_results(t_us, dVc_dt, y_a_opt, a_rounded):
        """绘制结果图表"""
        plt.figure(figsize=(8, 6))
        
        # 绘制曲线
        plt.plot(t_us, dVc_dt, label='dVc/dt', linewidth=1.5)
        plt.plot(t_us, y_a_opt, label=f'y_a (a/X={a_rounded})', 
                linestyle='--', linewidth=1.5)

        # 设置标签和字体
        plt.xlabel(r'Time [$\mu s$]', fontname='Arial', fontsize=12)
        plt.ylabel('Differential of $ \\left | V_c \\right |  $', 
                  fontname='Arial', fontsize=12)
        
        # 设置刻度
        plt.xticks(fontname='Arial', fontsize=12)
        plt.yticks(fontname='Arial', fontsize=12)
        
        # 添加网格和图例
        plt.grid(True, linestyle='--', linewidth=0.7)
        plt.legend(loc='best', fontsize=12, frameon=True, 
                  fancybox=True, edgecolor='black', framealpha=1)
        
        # 设置边框
        ax = plt.gca()
        for spine in ax.spines.values():
            spine.set_linewidth(1.5)
            
        plt.tight_layout()
        plt.show()

def process_file(filename):
    """处理单个数据文件"""
    # 读取原始数据
    raw_powers = read_hdf5_find_powers(filename, 2)
    
    try:
        voltages_IQ = mag_phase_to_IQ(raw_powers, Config.RQ, Config.Ql, Config.beta)
    except KeyError as e:
        print(f"处理文件 {filename} 时出错: {e}")
        return

    # 获取校准数据
    calibration_results = VoltageProcessor.calculate_calibration_matrix(voltages_IQ)
    if calibration_results is None:
        return

    # 准备计算参数
    t = np.arange(len(voltages_IQ["Vcav_I"])) * Config.Ts
    t_us = t * 1e6  # 转换为微秒

    # 提取复数形式的电压值
    Vf_star = voltages_IQ['Vfor_I'] + 1j * voltages_IQ['Vfor_Q']
    Vr_star = voltages_IQ['Vrefl_I'] + 1j * voltages_IQ['Vrefl_Q']
    Vc_star = voltages_IQ['Vcav_I'] + 1j * voltages_IQ['Vcav_Q']

    # 计算阻抗和优化参数
    Z = Cal_Z(Vf_star, Vr_star)
    a_opt = OptimizationCalculator.optimize_parameters(Vc_star, Vf_star, Vr_star, Z)

    # 计算和显示结果
    C = calibration_results['C_array']
    X = C[0] + 1j * C[1]
    Y = C[2] + 1j * C[3]
    
    a_round = a_opt / X
    a_rounded = complex(round(a_round.real, 2), round(a_round.imag, 2))
    
    print(f'a={a_opt:.3f}, b={a_opt/Z:.3f}, c={X-a_opt:.3f}, d={Y-a_opt/Z:.3f}')
    print(f'X={X:.2f}, Y={Y:.2f}, Z={Z:.2f}')

    # 计算最终结果并绘图
    dVc_dt = calculate_deri(np.abs(Vc_star), Config.Fs)
    y_a_opt = OptimizationCalculator.calculate_y_a(
        [a_opt.real, a_opt.imag], Z, Vc_star, Vf_star, Vr_star
    )
    
    Plotter.plot_results(t_us, dVc_dt, y_a_opt, a_rounded)

def main():
    """主函数"""
    for filename in DATA_SOURCES:
        process_file(filename)

if __name__ == '__main__':
    main()
